//
// Copyright 2019 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS-IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
//

import {LogicalCore, SystemTopology as ISystemTopology, SystemTopologyResponse} from '../models/collection_data_services';

import {CpuLabel, SystemTopology} from './system_topology';


/**
 * CpuLabel encapsulates the labeling for a single CPU in a possibly-multicore
 * system.  It is generated by a SystemTopology object, and knows its place
 * within that topology.
 * Version specific to internal architectures.
 */
export class ComplexCpuLabel extends CpuLabel {
  constructor(
      public logicalCore: LogicalCore,
      public systemTopology: ComplexSystemTopology) {
    super(Number(logicalCore.cpuId), Math.max(...systemTopology.cpus, 0));
  }

  /**
   * @return The hyperthread of this CPU.
   */
  getHyperThread(): number {
    return this.logicalCore.threadId;
  }

  /**
   * @return The NUMA node of this CPU
   */
  getNumaNode(): number {
    return this.logicalCore.numaNodeId;
  }

  /**
   * @return The core index within NUMA node.
   */
  getCore(): number {
    return this.logicalCore.coreId;
  }

  /**
   * Returns a label for this CPU.
   * @param includeHyperthread Include the hyperthread index?
   * @param changedComponentsOnly Only label transitions between
   *   hyperthreads, CPUs, NUMA?
   */
  getTopoLabel(includeHyperthread = false, changedComponentsOnly = false):
      string {
    let label = '';
    if (includeHyperthread) {
      label = `HT${this.getHyperThread()}`;
    }
    if (this.getHyperThread() === 0 || !changedComponentsOnly) {
      label += ` core ${this.getCore()}`;
    }
    if ((this.getHyperThread() === 0 && this.getCore() === 0) ||
        !changedComponentsOnly) {
      label += ` NUMA ${this.getNumaNode()}`;
    }
    return label;
  }
}

/**
 * Provides utilities for working with system topologies and CPUs,
 * including text filtering of CPU sets.
 * Version specific to internal architectures.
 */
export class ComplexSystemTopology extends SystemTopology {
  hyperthreadCountPerCore = 1;
  numaNodeCount = 1;
  blockSize = 1;
  maxCoreIDs: number[] = [];

  systemTopology: ISystemTopology;

  /**
   * Constructs a new SystemTopology from the provided SystemTopologyProto.
   * If the provided SystemTopologyProto is empty, it instead guesses a likely
   * topology from the provided list of CPU IDs.
   */
  constructor(response: SystemTopologyResponse, cpus: number[]) {
    super(response.systemTopology.logicalCores.map(cpuLabel => cpuLabel.cpuId));
    this.systemTopology = response.systemTopology;
    const validTopology = this.systemTopology &&
        this.systemTopology.logicalCores.length >= cpus.length;
    if (!validTopology) {
      if (this.systemTopology) {
        console.log(`Warning: Inconsistent system topology (${
            this.systemTopology.logicalCores.length} vs ${cpus.length} cores)`);
      }
      return;
    }

    this.cpuLabels = this.systemTopology.logicalCores.map(
        logicalCore => new ComplexCpuLabel(logicalCore, this));
    this.hyperthreadCountPerCore =
        this.cpuLabels
            .map((cpuLabel) => (cpuLabel as ComplexCpuLabel).getHyperThread())
            .reduce(
                (maxHyperthread, hyperthread) =>
                    (hyperthread > maxHyperthread) ? hyperthread :
                                                     maxHyperthread,
                0) +
        1;
    this.numaNodeCount =
        this.cpuLabels
            .map((cpuLabel) => (cpuLabel as ComplexCpuLabel).getNumaNode())
            .reduce(
                (maxNumaNode, numaNode) =>
                    (numaNode > maxNumaNode) ? numaNode : maxNumaNode,
                0) +
        1;

    this.maxCoreIDs = this.cpuLabels.reduce((maxCoreIDs, cpuLabel) => {
      const complexLabel = cpuLabel as ComplexCpuLabel;
      const numaNode = complexLabel.getNumaNode();
      const coreId = complexLabel.getCore();
      maxCoreIDs[numaNode] = Math.max(maxCoreIDs[numaNode], coreId);
      return maxCoreIDs;
    }, new Array(this.getNumaNodeCount()).fill(0));

    this.blockSize = Math.ceil(
        this.cpuCount / this.getHyperthreadCountPerCore() /
        this.getNumaNodeCount());

    (this.cpuLabels as ComplexCpuLabel[]).sort((a, b) => {
      const numaOrder = a.getNumaNode() - b.getNumaNode();
      if (numaOrder !== 0) {
        return numaOrder;
      }

      const coreOrder = a.getCore() - b.getCore();
      if (coreOrder !== 0) {
        return coreOrder;
      }

      const htOrder = a.getHyperThread() - b.getHyperThread();
      return htOrder;
    });
  }

  /**
   * @return The 'block size' of this topology: the number of physical cores per
   *   NUMA node.
   */
  getBlockSize(): number {
    return this.blockSize;
  }

  /**
   * @return The number of hyperthreads per core in this topology.
   */
  getHyperthreadCountPerCore(): number {
    return this.hyperthreadCountPerCore;
  }

  /**
   * @return The number of NUMA nodes in this topology.
   */
  getNumaNodeCount(): number {
    return this.numaNodeCount;
  }

  /**
   * @return The maximum core ID for each NUMA node
   */
  getMaxCoreIDs(): number[] {
    return this.maxCoreIDs;
  }

  /**
   * Takes a CPU selection string and returns an equivalent but canonicalized
   * one.
   */
  canonicalizeCpuFilterString(cpuSelectionString: string): string {
    // All others: remove whitespace.
    return super.canonicalizeCpuFilterString(cpuSelectionString);
  }

}
